import json

import pandas as pd
from xtools.utils import write_json_lines


def read_json_lines(file):
    return [json.loads(x) for x in open(file, encoding='utf-8')]


df = read_json_lines('d:/cluener_public/train.json')
test = read_json_lines('d:/cluener_public/dev.json')


def convert(data, out_name):
    out = []
    for line in data:
        label = line.pop('label')
        line['entities'] = []
        for k, v in label.items():
            tmp = {'entity_type': k, 'entity': next(iter(v.keys())), 'spans': next(iter(v.values()))[0]}
            line['entities'].append(tmp)
        out.append(line)

    write_json_lines(out, out_name)


convert(df, 'train.json')
convert(test, 'dev.json')